import os
import sys
import shutil
import subprocess
import platform
import site
import importlib
from pathlib import Path

# 获取项目根目录
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

# 虚拟环境路径
VENV_DIR = os.path.join(ROOT_DIR, "venv")

# 定义模型目录
MODEL_DIR = os.path.join(ROOT_DIR, "models")
SENSEVOICE_MODEL_DIR = os.path.join(MODEL_DIR, "sensevoice", "SenseVoiceSmall")
FSMN_VAD_MODEL_DIR = os.path.join(MODEL_DIR, "fsmn", "speech_fsmn_vad_zh-cn-16k-common-pytorch")

# 定义输出目录
DIST_DIR = os.path.join(ROOT_DIR, "dist")
BUILD_DIR = os.path.join(ROOT_DIR, "build")

# 定义图标文件路径（如果有）
ICON_PATH = os.path.join(ROOT_DIR, "icon.ico") if os.path.exists(os.path.join(ROOT_DIR, "icon.ico")) else None

def ensure_directory(path):
    """确保目录存在，如果不存在则创建"""
    if not os.path.exists(path):
        os.makedirs(path)

def check_missing_files(build_dir):
    """检查打包完成后是否缺少关键文件"""
    print("检查打包后的文件...")
    
    critical_files = [
        ("funasr/version.txt", "FunASR版本文件"),
    ]
    
    missing_files = []
    for file_path, description in critical_files:
        full_path = os.path.join(build_dir, "_internal", file_path)
        if not os.path.exists(full_path):
            missing_files.append((full_path, description))
    
    if missing_files:
        print("\n警告: 以下关键文件在打包后未找到:")
        for path, desc in missing_files:
            print(f"  - {desc}: {path}")
        print("这可能会导致应用程序无法正常运行。")
    else:
        print("所有关键文件已正确包含。")
    
    return missing_files

def check_models():
    """检查模型文件是否存在，如果不存在给出警告"""
    missing_models = []
    
    if not os.path.exists(SENSEVOICE_MODEL_DIR):
        missing_models.append(SENSEVOICE_MODEL_DIR)
    
    if not os.path.exists(FSMN_VAD_MODEL_DIR):
        missing_models.append(FSMN_VAD_MODEL_DIR)
    
    if missing_models:
        print("警告: 以下模型目录未找到:")
        for model in missing_models:
            print(f"  - {model}")
        print("请确保模型文件已下载并放置在正确的位置!")
        print("是否继续构建? (y/n)")
        response = input().strip().lower()
        if response != 'y':
            sys.exit(1)

def clean_previous_build():
    """清理之前的构建文件"""
    if os.path.exists(DIST_DIR):
        shutil.rmtree(DIST_DIR)
    if os.path.exists(BUILD_DIR):
        shutil.rmtree(BUILD_DIR)
    
    # 清理 __pycache__ 目录
    for root, dirs, files in os.walk(ROOT_DIR):
        if '__pycache__' in dirs:
            shutil.rmtree(os.path.join(root, '__pycache__'))

def find_package_files(package_name):
    """查找Python包的安装路径和数据文件"""
    print(f"尝试查找包: {package_name}")
    
    # 首先尝试从虚拟环境中查找
    venv_site_packages = None
    
    # 检查Windows路径
    if os.path.exists(VENV_DIR):
        # 查找Python版本目录
        lib_dir = os.path.join(VENV_DIR, "Lib")
        if os.path.exists(lib_dir):
            site_packages = os.path.join(lib_dir, "site-packages")
            if os.path.exists(site_packages):
                venv_site_packages = site_packages
                print(f"已找到虚拟环境site-packages: {site_packages}")
                
        # 查找Linux/Mac路径
        if not venv_site_packages:
            for root, dirs, files in os.walk(VENV_DIR):
                if os.path.basename(root) == "site-packages":
                    venv_site_packages = root
                    print(f"已找到虚拟环境site-packages: {venv_site_packages}")
                    break
    
    # 检查虚拟环境中的包
    if venv_site_packages:
        package_path = os.path.join(venv_site_packages, package_name)
        if os.path.exists(package_path):
            print(f"在虚拟环境中找到{package_name}包: {package_path}")
            return package_path
    
    # 如果在虚拟环境中没找到，尝试使用importlib
    try:
        package = importlib.import_module(package_name)
        print(f"成功导入{package_name}包")
        
        # 安全地获取文件路径
        try:
            if hasattr(package, '__file__'):
                package_path = os.path.dirname(package.__file__)
                print(f"使用__file__找到{package_name}包路径: {package_path}")
                return package_path
            else:
                print(f"警告: {package_name}包没有__file__属性")
                
                # 尝试使用__path__属性（对于包）
                if hasattr(package, '__path__') and len(package.__path__) > 0:
                    package_path = package.__path__[0]
                    print(f"使用__path__找到{package_name}包路径: {package_path}")
                    return package_path
                
                # 如果上述方法都失败，尝试在所有已知的site-packages中查找
                for path in site.getsitepackages():
                    potential_path = os.path.join(path, package_name)
                    if os.path.exists(potential_path):
                        print(f"在系统路径中找到{package_name}包: {potential_path}")
                        return potential_path
        except Exception as path_err:
            print(f"获取{package_name}包路径时出错: {str(path_err)}")
    except ImportError as e:
        print(f"无法导入{package_name}包: {str(e)}")
    except Exception as e:
        print(f"导入{package_name}包时发生未知错误: {str(e)}")
    
    # 如果所有方法都失败，返回None
    print(f"警告: 无法定位{package_name}包路径")
    return None

def build_executable():
    """构建可执行文件"""
    # 确保在虚拟环境中运行
    if VENV_DIR not in sys.prefix:
        print(f"警告: 当前不在虚拟环境中运行。为了确保正确打包，建议在虚拟环境中运行此脚本。")
        print(f"虚拟环境路径: {VENV_DIR}")
        print(f"当前Python路径: {sys.prefix}")
        print("是否继续? (y/n)")
        response = input().strip().lower()
        if response != 'y':
            sys.exit(1)
    
    # 构建 PyInstaller 命令行参数
    cmd = [
        "pyinstaller",
        "--name=SenseVoice2",
        "--windowed",  # GUI应用程序模式
        "--onedir",    # 单目录模式
        "--clean",     # 清理临时文件
        "--noconfirm", # 不确认覆盖
        "--log-level=WARN",  # 只显示警告和错误，减少日志输出
        "--collect-all=funasr",  # 收集所有funasr相关文件
    ]
    
    # 添加图标（如果有）
    if ICON_PATH:
        cmd.append(f"--icon={ICON_PATH}")
    
    # 添加主程序文件
    cmd.append("app.py")
    
    # 添加必要的隐式导入
    hidden_imports = [
        "torch",
        "torchaudio",
        "torchaudio.lib.libtorchaudio",  # 确保torchaudio库被正确包含
        "funasr",
        "funasr.auto",
        "funasr.auto.auto_model",  # 明确添加auto_model模块
        "funasr.register",         # 添加register模块，这是模型注册和加载的关键
        "funasr.models",           # 添加models模块
        "funasr.models.ctc.ctc",   # 添加CTC模型
        "funasr.utils",            # 添加utils模块
        "funasr.utils.postprocess_utils",  # 添加后处理工具
        "numpy",
        "PyQt5",
        "PyQt5.sip",  # 添加 PyQt5.sip 模块
        "sip",        # 添加 sip 模块
        "model",
        "tqdm",       # 添加tqdm模块
        "tqdm.std",   # 添加tqdm.std模块
    ]
    
    # 尝试加载 scipy 并添加相关依赖
    try:
        import scipy
        hidden_imports.append("scipy")
        hidden_imports.append("scipy.special")
        hidden_imports.append("scipy.special._ufuncs")
        hidden_imports.append("scipy.special._ufuncs_cxx")
        hidden_imports.append("scipy.special._ellip_harm_2")
        hidden_imports.append("scipy.special._cdflib")
        print("检测到scipy，添加相关依赖")
    except ImportError:
        print("未检测到scipy，跳过相关依赖")
    
    for imp in hidden_imports:
        cmd.append(f"--hidden-import={imp}")
    
    
    # 添加数据文件
    cmd.append(f"--add-data=model.py{os.pathsep}.")
    cmd.append(f"--add-data=style.qss{os.pathsep}.")
    cmd.append(f"--add-data=ai.ico{os.pathsep}.")
    
    # 添加FunASR库的数据文件
    funasr_path = find_package_files("funasr")
    if funasr_path:
        # 添加version.txt文件
        version_file = os.path.join(funasr_path, "version.txt")
        if os.path.exists(version_file):
            print(f"添加FunASR版本文件: {version_file}")
            cmd.append(f"--add-data={version_file}{os.pathsep}funasr")
        else:
            print(f"警告: FunASR版本文件未找到: {version_file}")
            # 创建一个临时版本文件
            temp_version_file = os.path.join(ROOT_DIR, "funasr_version.txt")
            with open(temp_version_file, "w") as f:
                f.write("1.1.3\n")  # 使用requirements.txt中指定的版本
            cmd.append(f"--add-data={temp_version_file}{os.pathsep}funasr")
            
        # 添加其他可能的数据目录
        for dirname in ["configs", "bin", "data"]:
            dir_path = os.path.join(funasr_path, dirname)
            if os.path.exists(dir_path) and os.path.isdir(dir_path):
                print(f"添加FunASR数据目录: {dir_path}")
                cmd.append(f"--add-data={dir_path}{os.pathsep}funasr/{dirname}")
                
        # 确保所有Python包路径都被正确收集
        cmd.append(f"--paths={funasr_path}")
        
        # 特别处理 - 创建一个空的version.txt文件在funasr包内部
        # 这是为了确保PyInstaller能找到它
        if not os.path.exists(version_file):
            os.makedirs(os.path.dirname(version_file), exist_ok=True)
            with open(version_file, "w") as f:
                f.write("1.1.3\n")
            print(f"已在源代码位置创建临时版本文件: {version_file}")
            
    # 添加虚拟环境的site-packages路径
    if os.path.exists(VENV_DIR):
        # Windows路径
        site_packages = os.path.join(VENV_DIR, "Lib", "site-packages")
        if os.path.exists(site_packages):
            cmd.append(f"--paths={site_packages}")
        
        # Linux/Mac路径
        for root, dirs, files in os.walk(VENV_DIR):
            if os.path.basename(root) == "site-packages":
                cmd.append(f"--paths={root}")
                break
    
    # 添加需要排除的模块（可选）
    excludes = [
        "matplotlib",
        "IPython",
        "jupyter_client",
        "jupyter_core",
        "notebook",
        "qtconsole",
        "torch.utils.tensorboard",  # 排除 tensorboard 相关模块，解决警告
        "tensorboard"
    ]
    for ex in excludes:
        cmd.append(f"--exclude-module={ex}")
    
    # 添加模型目录作为数据文件（如果存在）
    if os.path.exists(MODEL_DIR):
        cmd.append(f"--add-data={MODEL_DIR}{os.pathsep}models")
    
    # 设置其他 PyInstaller 选项
    cmd.append("--noupx")  # 不使用 UPX 压缩
    
    # 确保输出目录存在
    ensure_directory("mydata/output")
    cmd.append(f"--add-data=mydata{os.pathsep}mydata")
    
    # 打印构建命令
    print("执行构建命令:", " ".join(cmd))
    
    # 执行 PyInstaller 命令
    subprocess.run(cmd, check=True)

def copy_additional_files():
    """复制额外文件到构建目录"""
    app_dir = os.path.join(DIST_DIR, "SenseVoice2")
    
    # 创建输出目录（如果还不存在）
    output_dir = os.path.join(app_dir, "mydata", "output")
    ensure_directory(output_dir)
    
    # 检查funasr版本文件是否存在，如果不存在则创建
    funasr_version_file = os.path.join(app_dir, "_internal", "funasr", "version.txt")
    if not os.path.exists(funasr_version_file):
        print(f"在打包后添加缺失的FunASR版本文件: {funasr_version_file}")
        os.makedirs(os.path.dirname(funasr_version_file), exist_ok=True)
        with open(funasr_version_file, "w") as f:
            f.write("1.1.3\n")  # 使用requirements.txt中指定的版本
    
    # 创建 README 文件（可选）
    with open(os.path.join(app_dir, "README.txt"), "w", encoding="utf-8") as f:
        f.write("SenseVoice2\n")
        f.write("======================\n\n")
        f.write("使用说明：\n")
        f.write("1. 点击\"上传音频文件\"按钮上传音频文件\n")
        f.write("2. 点击\"开始转换\"按钮进行语音识别\n")
        f.write("3. 转换结果将显示在下方文本框，也可保存到文件\n")
        f.write("4. 所有转换结果默认保存在程序目录下的 mydata/output 文件夹\n\n")
        f.write("注意事项：\n")
        f.write("- 支持格式：wav, mp4, m4a\n")
            
        f.write("- 首次启动时可能需要等待模型加载\n")

def handle_warnings():
    """处理常见的PyInstaller警告"""
    print("提示: 在构建过程中可能会出现一些警告，大多数可以安全忽略，例如:")
    print("  - 关于'torch.utils.tensorboard'的警告 (您的应用不需要使用TensorBoard)")
    print("  - 关于'sip'模块的警告 (这是PyQt5的一部分，但可能不影响GUI功能)")
    print("  - 关于'scipy.special._cdflib'的警告 (如果您的应用不使用scipy特殊函数可忽略)")
    print("  - 模块收集警告 (不影响主要功能的模块)")
    print("  - 未找到某些DLL的警告 (通常是因为您的系统上没有这些可选组件)")
    print("只要最终生成了可执行文件，这些警告通常不会影响应用功能。\n")

def main():
    """主函数，执行构建流程"""
    print("开始构建 SenseVoice 应用...")
    
    # 检查环境
    print(f"操作系统: {platform.platform()}")
    print(f"Python版本: {platform.python_version()}")
    
    # 输出警告处理信息
    handle_warnings()
    
    # 确保必要的目录存在
    ensure_directory("mydata/output")
    
    # 检查模型文件
    check_models()
    
    # 清理之前的构建
    print("清理之前的构建文件...")
    clean_previous_build()
    
    # 构建可执行文件
    print("开始构建可执行文件...")
    build_executable()
    
    # 复制额外文件
    print("复制额外文件...")
    copy_additional_files()
    
    # 检查是否缺少关键文件
    missing_files = check_missing_files(os.path.join(DIST_DIR, "SenseVoice2"))
    
    print("\n构建完成!")
    print(f"可执行文件已创建在: {os.path.join(DIST_DIR, 'SenseVoice2')}")
    print("注意: 如果您在运行应用时遇到问题，可能需要安装Visual C++ Redistributable包")

if __name__ == "__main__":
    main()
